import os
import argparse
from core.deploy import MiniCPMVDeploy
from inference.image_inference import ImageInferencer
from utils.batch_processor import process_directory


def main():
    """
    批量图片推理主函数
    """
    print("===== MiniCPM-V 批量图片推理工具 ======")
    print("基于模块化架构的批量图片处理")
    print("====================================")
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='MiniCPM-V 批量图片推理工具')
    # 设置默认图片目录为项目中的images文件夹
    default_image_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'images')
    parser.add_argument('--image_dir', type=str, default=default_image_dir,
                        help=f'图片文件夹路径，默认为: {default_image_dir}')
    parser.add_argument('--prompt', type=str, default=None,
                        help='提示词模板，支持{file_name}变量')
    parser.add_argument('--output_dir', type=str, default=None,
                        help='输出结果目录，默认为配置文件中的output_dir')
    parser.add_argument('--export_csv', type=str, default=None,
                        help='导出CSV文件路径')
    parser.add_argument('--max_workers', type=int, default=None,
                        help='最大工作线程数，默认根据GPU内存自动计算')
    parser.add_argument('--recursive', action='store_true',
                        help='是否递归搜索子目录')
    parser.add_argument('--no_save', action='store_true',
                        help='不保存单个结果文件')
    
    args = parser.parse_args()
    
    # 验证图片目录是否存在
    if not os.path.exists(args.image_dir):
        print(f"错误: 默认目录不存在: {args.image_dir}")
        print("请输入图片文件夹路径: ")
        args.image_dir = input().strip()
        if not args.image_dir or not os.path.exists(args.image_dir):
            print("未提供有效的图片文件夹路径，程序退出")
            return
    
    # 检查目录是否存在
    if not os.path.exists(args.image_dir):
        print(f"错误: 目录不存在: {args.image_dir}")
        return
    
    # 如果没有提供提示词，使用默认提示词模板
    if not args.prompt:
        print("请输入提示词模板（留空使用默认模板）: ")
        args.prompt = input().strip() or "请详细描述这张图片的内容"
    
    # 显示参数信息
    print("\n===== 批量推理参数 =====")
    print(f"图片目录: {args.image_dir}")
    print(f"提示词模板: {args.prompt}")
    print(f"是否递归搜索: {args.recursive}")
    if args.max_workers is not None:
        print(f"最大工作线程数: {args.max_workers}")
    if args.export_csv:
        print(f"CSV导出路径: {args.export_csv}")
    print(f"是否保存单个结果文件: {not args.no_save}")
    print("====================")
    
    # 初始化部署类
    deploy = MiniCPMVDeploy()
    
    # 更新输出目录配置（如果提供）
    if args.output_dir:
        deploy.update_config({"paths": {"output_dir": args.output_dir}})
    
    # 开始批量处理
    process_directory(
        dir_path=args.image_dir,
        inferencer_class=ImageInferencer,
        prompt=args.prompt,
        save_result=not args.no_save,
        export_csv=args.export_csv,
        max_workers=args.max_workers,
        recursive=args.recursive,
        deploy=deploy
    )
    
    # 卸载模型，释放内存
    deploy.unload_model()
    
    print("\n===== 批量图片推理完成 ======")
    print("感谢使用MiniCPM-V批量图片推理工具！")


if __name__ == "__main__":
    main()